/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

#include <getopt.h>
#include "enclave.h"
#include "util.h"
#define  QT_ENCLAVE_QUERY_ALL                -1
#define  QT_ENCLAVE_MAX                      2
#define  QT_MIN_ENCLAVE_ID                   0
#define  QT_MAX_ENCLAVE_ID                   1

static void help_build_image(void)
{
    printf("%s build_image -k vmlinuz -i initramfs [-c cmdline]\n"
           "[-p private_key] [-s signing_certificate] -o output\n", ENCLAVE_BIN);
}

static void help_enclave_run(void)
{
    printf("%s run [-c cpus(default 2)] [-m memory(MB default 1024)]\n"
           "-i image [-s cid (default 4)]\n", ENCLAVE_BIN);
}

static void help_describe_single_enclave(void)
{
    printf("%s describe [-i enclave_id]\n", ENCLAVE_BIN);
}

static void help_enclave_destroy(void)
{
    printf("%s destroy -i enclave_id\n", ENCLAVE_BIN);
}

static void help_enclave_console(void)
{
    printf("%s console -i enclave_id \n", ENCLAVE_BIN);
}

static int enclave_info_check(struct enclave_instance_info *enclave_info)
{
    if (!enclave_info) {
        help_enclave_run();
        return -1;
    }

    if (!enclave_info->image) {
        qt_log_err(QT_MISSING_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Missing image parameters!");
        return -1;
    }

    if (!enclave_info->cpus) {
        qt_log_err(QT_MISSING_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Missing CPUs parameters!");
        return -1;
    }

    if (!enclave_info->memory) {
        qt_log_err(QT_MISSING_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Missing memory parameters!");
        return -1;
    }

    if (!enclave_info->vsock_cid) {
        qt_log_err(QT_MISSING_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Missing cid parameters!");
        return -1;
    }

    return 0;
}

static int cli_enclave_run(int argc, char **argv)
{
    int c;
    int64_t value;
    struct enclave_instance_info enclave_info = { 0 };

    const struct option options[] = {
        { "help", no_argument, 0, 'h' },
        { "cpus", required_argument, 0, 'c' },
        { "memory", required_argument, 0, 'm' },
        { "image", required_argument, 0, 'i' },
        { "cid", required_argument, 0, 's' },
        { "debug", no_argument, 0, 'd' },
        { 0, 0, 0, 0 }
    };

    while ((c = getopt_long(argc, argv, "hc:m:i:s:d", options,
                            NULL)) != -1) {
        switch (c) {
            case 'h':
            case '?':
                help_enclave_run();
                return 0;
            case 'c':
                value = atoll(optarg);
                if (value > 0 && value <= INT_MAX)
                    enclave_info.cpus = (int)value;
                else {
                    qt_log_err(QT_INVALID_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                               "Invalid cpus: %d", value);
                    return -1;
                }
                break;
            case 'm':
                value = atoll(optarg);
                if (value > 0)
                    enclave_info.memory = value;
                else {
                    qt_log_err(QT_INVALID_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                               "Invalid memory: %d", value);
                    return -1;
                }
                break;
            case 'i':
                enclave_info.image = optarg;
                break;
            case 's':
                value = atoll(optarg);
                if (value > 0)
                    enclave_info.vsock_cid = (int64_t)value;
                else {
                    qt_log_err(QT_INVALID_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                               "Invalid cid: %d", value);
                    return -1;
                }
                break;
            case 'd':
                enclave_info.debug = 1;
                break;
            default:
                printf("unsupport option %c\n", c);
                break;
        }
    }

    if (enclave_info_check(&enclave_info) < 0)
        return -1;

    return enclave_run(&enclave_info);
}

static int cli_describe_single_enclave(int argc, char **argv)
{
    int c;
    int64_t enclave_id = -1;

    const struct option options[] = {
        {"help", no_argument, 0, 'h'},
        {"enclave_id", optional_argument, 0, 'i'},
        {0, 0, 0, 0}
    };

    while ((c = getopt_long(argc, argv, "hi::", options, NULL)) != -1) {
        switch (c) {
            case 'h':
            case '?':
                help_describe_single_enclave();
                return 0;
            case 'i':
                enclave_id = atoll(optarg);
                break;
            default:
                printf("unsupport option %c\n", c);
                break;
        }
    }

    if (enclave_id < 0 || enclave_id >= QT_ENCLAVE_MAX) {
        qt_log_err(QT_INVALID_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "MIN enclave-id: %d, MAX enclave-id: %d",
                   QT_MIN_ENCLAVE_ID, QT_MAX_ENCLAVE_ID);
        return -1;
    }

    return enclave_describes(enclave_id);
}

static int cli_describe_all_enclaves()
{
    int rc;
    rc = enclave_describes(QT_ENCLAVE_QUERY_ALL);
    if (rc < 0) {
        qt_log_err(QT_UNDEINED_ERROR_NUM, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Describe all enclaves failure");
    }
    return rc;
}

static int cli_enclave_destroy(int argc, char **argv)
{
    int c;
    int64_t enclave_id = -1;

    const struct option options[] = {
        {"help", no_argument, 0, 'h'},
        {"enclave_id", required_argument, 0, 'i'},
        {0, 0, 0, 0}
    };

    while ((c = getopt_long(argc, argv, "hi:",
                            options, NULL)) != -1) {
        switch (c) {
            case 'h':
            case '?':
                help_enclave_destroy();
                return 0;
            case 'i':
                enclave_id = atoll(optarg);
                break;
            default:
                printf("unsupport option %c\n", c);
                break;
        }
    }

    if (enclave_id < 0) {
        qt_log_err(QT_INVALID_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "MIN enclave-id: %d, MAX enclave-id: %d",
                   QT_MIN_ENCLAVE_ID, QT_MAX_ENCLAVE_ID);
        return -1;
    }

    return enclave_destroy(enclave_id);
}

static int cli_enclave_console(int argc, char **argv)
{
    int c;
    int64_t enclave_id = -1;

    const struct option options[] = {
        {"help", no_argument, 0, 'h'},
        {"enclave_id", optional_argument, 0, 'i'},
        {0, 0, 0, 0}
    };

    while ((c = getopt_long(argc, argv, "hi:",
                            options, NULL)) != -1) {
        switch (c) {
            case 'h':
            case '?':
                help_enclave_console();
                return 0;
            case 'i':
                enclave_id = atoll(optarg);
                break;
            default:
                printf("unsupport option %c \n", c);
                break;
        }
    }

    if (enclave_id < 0) {
        qt_log_err(QT_INVALID_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "MIN enclave-id: %d, MAX enclave-id: %d",
                   QT_MIN_ENCLAVE_ID, QT_MAX_ENCLAVE_ID);
        return -1;
    }

    return enclave_console(enclave_id);
}

static int cli_build_image(int argc, char **argv)
{
    int c;
    struct enclave_img_info img_info = { 0 };
    const struct option options[] = {
        {"help", no_argument, 0, 'h'},
        {"kernel", required_argument, 0, 'k'},
        {"initrd", required_argument, 0, 'i'},
        {"cmdline", required_argument, 0, 'c'},
        {"output", required_argument, 0, 'o'},
        {"private_key", required_argument, 0, 'p'},
        {"signing_certificate", required_argument, 0, 's'},
        {0, 0, 0, 0}
    };

    while ((c = getopt_long(argc, argv, "hk:i:c:o:p:s:",
                            options, NULL)) != -1) {
        switch (c) {
            case 'h':
            case '?':
                help_build_image();
                return 0;
            case 'k':
                img_info.kernel = optarg;
                break;
            case 'i':
                img_info.initrd = optarg;
                break;
            case 'c':
                img_info.cmdline = optarg;
                break;
            case 'o':
                img_info.output = optarg;
                break;
            case 'p':
                img_info.private_key = optarg;
                break;
            case 's':
                img_info.signing_certificate = optarg;
                break;
            default:
                printf("unsupport option %c\n", c);
                break;
        }
    }

    if (!img_info.initrd || !img_info.kernel || !img_info.output) {
        help_build_image();
        return -1;
    }

    if ((img_info.private_key && !img_info.signing_certificate) ||
        (!img_info.private_key && img_info.signing_certificate)) {
        printf("please specify private-key and signing-certificate at the\n"
               "same time\n");
        return -1;
    }

    return enclave_img_build(&img_info);
}

static void help(void)
{
    printf("Usage: %s command [command options]\n\n", ENCLAVE_BIN);
    help_build_image();
    help_enclave_run();
    help_enclave_destroy();
    help_describe_single_enclave();
}

struct enclave_cmd {
    const char *name;
    int (*handler)(int argc, char **argv);
};

static struct enclave_cmd enclave_cmds[] = {
    { "build_image", cli_build_image },
    { "run", cli_enclave_run },
    { "destroy", cli_enclave_destroy },
    { "describe_single_enclave", cli_describe_single_enclave},
    { "describe_all_enclaves", cli_describe_all_enclaves},
    { "console", cli_enclave_console },
    { NULL, NULL }
};

int main(int argc, char **argv)
{
	int rc;
	int c;
	struct enclave_cmd *cmd;
	const char *cmdname;
	static struct option options[] = {
		{ "help", no_argument, 0, 'h' },
		{ 0, 0, 0, 0 }
	};

	rc = logger_init();
	if (rc < 0) {
		printf("qt enclave:logger init err\n");
		return rc;
	}
	while ((c = getopt_long(argc, argv, "+h", options, NULL)) != -1) {
		if (c  == 'h' || c  == '?') {
			help();
			rc = 0;
			goto del_logger;
		}
	}

	cmdname = argv[optind];

	argc -= optind;
	if (argc < 1) {
		help();
		rc = -1;
		goto del_logger;
	}

	argv += optind;
	optind = 0;

	for (cmd = enclave_cmds; cmd->name != NULL; cmd++) {
		if (!strcmp(cmdname, cmd->name)) {
			rc = cmd->handler(argc, argv);
			goto del_logger;
		}
	}

	help();
del_logger:
	logger_del();
	return rc;
}
